#!/usr/bin/env python3

import argparse
import getpass
import ipaddress
import json
import ssl
import sys
import urllib.request
import urllib.parse
from pyVim.connect import SmartConnect, Disconnect
from pyVmomi import vim

# ---------- 参数 ----------
def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--vc-host", required=True)
    p.add_argument("--vc-user", required=True)
    p.add_argument("--vc-pass")
    p.add_argument("--tech-url", default="http://192.168.68.110:5380")
    p.add_argument("--token",   required=True,
                   help="Technitium API token（一次性登录后拿到）")
    p.add_argument("--zone",    required=True,
                   help="Zone 名，例如 hello.example")
    return p.parse_args()

# ---------- vCenter ----------
def get_vms(host, user, pwd):
    ctx = ssl._create_unverified_context()
    si = SmartConnect(host=host, user=user, pwd=pwd, sslContext=ctx)
    try:
        view = si.content.viewManager.CreateContainerView(
            si.content.rootFolder, [vim.VirtualMachine], True)
        vms = {}
        for vm in view.view:
            if vm.runtime.powerState != "poweredOn":
                continue
            ip = vm.guest.ipAddress
            if not ip:
                continue
            try:
                ipaddress.ip_address(ip)
            except ValueError:
                continue
            name = vm.name.replace(" ", "-").lower()
            vms[name] = ip
    finally:
        Disconnect(si)
    return vms

# ---------- Technitium ----------
def tech_add_a(args, host, ip):
    fqdn = f"{host}.{args.zone}".strip(".")
    params = {
        "token": args.token,
        "zone": args.zone,
        "domain": fqdn,
        "type": "A",
        "ipAddress": ip,
        "overwrite": "true",      # 允许覆盖
        "ttl": "3600",
        # 其余字段保持默认
    }
    url = f"{args.tech_url.rstrip('/')}/api/zones/records/add?" + urllib.parse.urlencode(params)
    with urllib.request.urlopen(url) as resp:
        data = json.loads(resp.read())
    if data.get("status") != "ok":
        print(f"⚠️ 失败 {fqdn}: {data}")
    else:
        print(f"✅ 添加/更新 {fqdn} -> {ip}")

# ---------- 主流程 ----------
def main():
    args = parse_args()
    vc_pass = args.vc_pass or getpass.getpass("vCenter 密码: ")
    vms = get_vms(args.vc_host, args.vc_user, vc_pass)

    for vm, ip in vms.items():
        tech_add_a(args, vm, ip)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit("\nCanceled")
    except Exception as e:
        sys.exit(f"❌ 错误: {e}")